package dslab.util;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import java.io.*;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

public class EncryptedReader extends BufferedReader {
    private final Cipher cipher;
    private boolean initialized = false;

    public EncryptedReader(InputStreamReader reader) throws NoSuchPaddingException, NoSuchAlgorithmException {
        super(reader);
        this.cipher = Cipher.getInstance("AES/CTR/NoPadding");
    }

    public String readUnsafeLine() throws IOException, IllegalBlockSizeException, BadPaddingException {
        String string = super.readLine();
        if (initialized)
            string = CryptoUtil.decryptMessage(string, cipher);
        return string;
    }

    public String readLine() throws IOException {
        String string = super.readLine();
        if (initialized) {
            try {
                string = CryptoUtil.decryptMessage(string, cipher);
            } catch (IllegalBlockSizeException | BadPaddingException exception) {
                throw new IOException(exception);
            }
        }
        return string;
    }

    public void init(SecretKey key, byte[] iv) throws InvalidAlgorithmParameterException, InvalidKeyException {
        cipher.init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(iv));
        initialized = true;
    }
}
